package com.own.component.elastic.service.config;

import co.elastic.clients.elasticsearch._types.SortOptions;
import co.elastic.clients.elasticsearch._types.query_dsl.BoolQuery;
import co.elastic.clients.elasticsearch._types.query_dsl.FunctionBoostMode;
import co.elastic.clients.elasticsearch._types.query_dsl.FunctionScore;
import co.elastic.clients.elasticsearch._types.query_dsl.FunctionScoreMode;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch.core.CountRequest;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.search.Highlight;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.elasticsearch.core.search.SourceConfig;
import co.elastic.clients.util.ObjectBuilder;
import com.alibaba.fastjson2.JSON;
import com.own.component.common.base.entity.po.BaseBean;
import com.own.component.elastic.entity.BaseElasticPo;
import com.own.component.elastic.entity.BaseElasticQuery;
import com.own.component.elastic.entity.BaseElasticVo;
import com.own.component.elastic.model.CustomFieldValue;
import lombok.Builder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * ElasticSearchConfig
 *
 * @author chenxueli
 * @date 2023-11-21 11:06:00
 */
@Slf4j
@Builder
public class ElasticSearchConfig<
        PO extends BaseBean,
        ELASTIC_PO extends BaseElasticPo<PO>,
        ELASTIC_VO extends BaseElasticVo<ELASTIC_PO>,
        QUERY extends BaseElasticQuery
        > {

    private static final Function<SourceConfig.Builder, ObjectBuilder<SourceConfig>> NULL_SOURCE = x -> x.filter(y -> y.includes(List.of()));
    private static final Consumer<BoolQuery.Builder> EMPTY_CONSUMER = x -> {
    };

    /**
     * 顶层的条件信息构造器（优先级更高）
     */
    @Builder.Default
    private BiFunction<QUERY, Consumer<BoolQuery.Builder>, Function<Query.Builder, ObjectBuilder<Query>>> queryBuilder = null;

    /**
     * 布尔查询条件信息构造器（低优先级）
     */
    @Builder.Default
    private BiConsumer<QUERY, BoolQuery.Builder> boolQueryBuilder = null;

    /**
     * 返回的内容
     */
    @Builder.Default
    private Function<SourceConfig.Builder, ObjectBuilder<SourceConfig>> source = null;

    /**
     * 排序信息构造器（单个）
     */
    @Builder.Default
    private Function<QUERY, Function<SortOptions.Builder, ObjectBuilder<SortOptions>>> sortBuilder = null;

    /**
     * 排序信息构造器（多个）
     */
    @Builder.Default
    private Function<QUERY, List<Function<SortOptions.Builder, ObjectBuilder<SortOptions>>>> sortBuilders = null;

    /**
     * 自定义分数权重构造器（单个）
     */
    @Builder.Default
    private Function<QUERY, Function<FunctionScore.Builder, ObjectBuilder<FunctionScore>>> functionScoreBuilder = null;

    /**
     * 自定义分数权重构造器（多个）
     */
    @Builder.Default
    private Function<QUERY, List<Function<FunctionScore.Builder, ObjectBuilder<FunctionScore>>>> functionScoreBuilders = null;

    /**
     * 高亮信息构造器
     */
    @Builder.Default
    private Function<QUERY, ElasticSearchHighlightHandler> highlightHandler = null;

    /**
     * hit处理器
     */
    @Builder.Default
    private Function<Hit<ELASTIC_PO>, ELASTIC_VO> hitHandler = null;

    /**
     * function分数权重模式信息
     */
    @Builder.Default
    private FunctionScoreMode functionScoreMode = FunctionScoreMode.Sum;

    /**
     * function分数权重模式信息
     */
    @Builder.Default
    private FunctionBoostMode functionBoostMode = FunctionBoostMode.Multiply;

    /**
     * 获取最终的查询条件信息构造器
     *
     * @param query    查询条件
     * @param consumer 其余自定义的筛选
     * @return 列表信息
     */
    public Function<Query.Builder, ObjectBuilder<Query>> getResultQueryBuilder(QUERY query, Consumer<BoolQuery.Builder> consumer) {
        // 如果存在顶层构造器
        if (queryBuilder != null) {
            return consumer == null ? queryBuilder.apply(query, EMPTY_CONSUMER) : queryBuilder.apply(query, consumer);
        }
        // 其余情况使用布尔查询构造器
        if (boolQueryBuilder != null) {
            return builder -> builder.bool(BoolQuery.of(b -> {
                boolQueryBuilder.accept(query, b);
                // 如果存在自定义的筛选条件
                if (consumer != null) {
                    consumer.accept(b);
                }
                return b;
            }));
        }
        // 如果存在自定义的筛选条件
        if (consumer != null) {
            return builder -> builder.bool(BoolQuery.of(b -> {
                consumer.accept(b);
                return b;
            }));
        }
        return null;
    }

    /**
     * 获取最终的排序信息构造器
     *
     * @param query 查询条件
     * @return 列表信息
     */
    public List<Function<SortOptions.Builder, ObjectBuilder<SortOptions>>> getResultSortBuilders(QUERY query) {
        return resolveMerge(query, sortBuilder, sortBuilders);
    }

    /**
     * 获取最终的自定义分数权重构造器
     *
     * @param query 查询条件
     * @return 列表信息
     */
    public List<Function<FunctionScore.Builder, ObjectBuilder<FunctionScore>>> getResultFunctionScoreBuilders(QUERY query) {
        return resolveMerge(query, functionScoreBuilder, functionScoreBuilders);
    }

    /**
     * 获取最终的高亮信息构造器
     *
     * @param query 查询条件
     * @return 信息
     */
    public Function<Highlight.Builder, ObjectBuilder<Highlight>> getResultHighlightBuilder(QUERY query) {
        if (highlightHandler == null) {
            return null;
        }
        return highlightHandler.apply(query).init();
    }

    /**
     * 获取最终的hit处理器
     *
     * @return 信息
     */
    public Function<Hit<ELASTIC_PO>, ELASTIC_VO> getResultHitHandler() {
        return hitHandler;
    }

    /**
     * 构造查询的基础信息
     *
     * @param index    文档名称
     * @param query    查询条件
     * @param builder  查询构造器
     * @param consumer 其余自定义的筛选
     * @return 查询构造器
     */
    public SearchRequest.Builder search(String index, QUERY query, Consumer<BoolQuery.Builder> consumer, SearchRequest.Builder builder) {
        // 设置基础信息
        builder.index(index)
                .from(StringUtils.isNotBlank(query.getSearchAfter()) ? 0 : query.getFrom())
                .size(query.getRows());
        if (source != null) {
            builder.source(source);
        }
        // 获取构造器信息
        var queryBuilder = getResultQueryBuilder(query, consumer);
        var functionScoreBuilders = getResultFunctionScoreBuilders(query);
        var sortBuilders = getResultSortBuilders(query);
        var highlightBuilder = getResultHighlightBuilder(query);
        // 分数权重构造器
        if (functionScoreBuilders != null && !functionScoreBuilders.isEmpty()) {
            builder.query(b -> b.functionScore(fun -> {
                functionScoreBuilders.forEach(fun::functions);
                fun.boostMode(functionBoostMode);
                fun.scoreMode(functionScoreMode);
                // 搜索条件
                if (queryBuilder != null) {
                    fun.query(queryBuilder);
                }
                return fun;
            }));
        } else {
            // 搜索条件
            if (queryBuilder != null) {
                builder.query(queryBuilder);
            }
        }
        // 排序条件
        if (sortBuilders != null) {
            sortBuilders.forEach(builder::sort);
        }
        // 高亮操作（查询查询信息）
        if (highlightBuilder != null) {
            builder.highlight(highlightBuilder);
        }
        // 加入上一条的搜索记录
        if (StringUtils.isNotBlank(query.getSearchAfter())) {
            var searchAfter = JSON.parseArray(query.getSearchAfter(), CustomFieldValue.class)
                    .stream().map(CustomFieldValue::convert)
                    .collect(Collectors.toList());
            builder.searchAfter(searchAfter);
        }
        return builder;
    }

    /**
     * 构造聚合查询的基础信息
     *
     * @param index          文档名称
     * @param query          查询条件
     * @param builder        查询构造器
     * @param filterConsumer 其余自定义的筛选
     * @return 聚合查询的构造器
     */
    public SearchRequest.Builder aggregation(String index, QUERY query, Consumer<BoolQuery.Builder> filterConsumer, SearchRequest.Builder builder) {
        // 设置基础信息
        builder.index(index).size(0).source(NULL_SOURCE);
        // 获取构造器信息
        var queryBuilder = getResultQueryBuilder(query, filterConsumer);
        // 搜索条件
        if (queryBuilder != null) {
            builder.query(queryBuilder);
        }
        return builder;
    }

    /**
     * 构造数量查询的基础信息
     *
     * @param index          文档名称
     * @param query          查询条件
     * @param builder        查询构造器
     * @param filterConsumer 其余自定义的筛选
     * @return 查询构造器
     */
    public CountRequest.Builder count(String index, QUERY query, Consumer<BoolQuery.Builder> filterConsumer, CountRequest.Builder builder) {
        builder.index(index);
        var queryBuilder = getResultQueryBuilder(query, filterConsumer);
        if (queryBuilder != null) {
            builder.query(queryBuilder);
        }
        return builder;
    }

    /**
     * 合并信息
     *
     * @param query     查询条件
     * @param function  函数
     * @param functions 函数列表
     * @param <T>       泛型
     * @return 列表信息
     */
    private <T> List<T> resolveMerge(
            QUERY query,
            Function<QUERY, T> function,
            Function<QUERY, List<T>> functions
    ) {
        var list = new ArrayList<T>();
        if (function != null) {
            var item = function.apply(query);
            if (item != null) {
                list.add(item);
            }
        }
        if (functions != null) {
            var itemList = functions.apply(query);
            if (itemList != null) {
                list.addAll(itemList);
            }
        }
        return list;
    }

}
